import pandas as pd
import numpy as np
from matplotlib import pyplot as plt
import os

def draw2DFigure(probe_data_obj, figure_conf):
    img_dict = {}
    for key in probe_data_obj:
        # print("二维彩图：", key, probe_data_obj[key])
        df = probe_data_obj[key].pivot(index='Height',columns='DataStartTime',values='GumResult')
        fig, ax = plt.subplots()
        X, Y = np.meshgrid(np.arange(len(df.columns)), df.index)
        # print("二维彩图df：", df)
        c = ax.contour(X, Y, df, levels=8, colors="black", linewidths=0.5)
        ax.clabel(c, inline=True, fontsize=10)
        c = ax.contourf(X, Y, df, levels=40, cmap='jet')
        plt.colorbar(c)
        ax.set_xlabel('时间')
        ax.set_ylabel('高度(km)')
        if key == "T":
            ax.set_title('温度测量不确定度时空分布(K)')
        elif key == "D":
            ax.set_title('大气密度测量不确定度时空分布(%)')
        elif key == "N":
            ax.set_title('钠密度测量不确定度时空分布(%)')
        elif key == "F":
            ax.set_title('铁密度测量不确定度时空分布(%)')
        elif key == "W1":
            ax.set_title('纬向风测量不确定度时空分布(m/s)')
        elif key == "W2":
            ax.set_title('经向风测量不确定度时空分布(m/s)')
        elif key == "W":
            ax.set_title('风速测量不确定度时空分布(m/s)')
        elif key == "X":
            ax.set_title('风向测量不确定度时空分布(°)')
        fig_path = os.path.join(figure_conf["figure_location"], key+'.png')
        fig.savefig(fig_path, dpi=300, bbox_inches="tight")
        img_dict[key] = fig_path
    return img_dict

# def dataCompareDrawIndexFigure(probe_data_obj, figure_conf):
#     return img_dict

def dataAnalysisDrawSpaceTime(probe_data_obj, figure_conf):
    img_dict = {}
    for key in probe_data_obj:
        df = probe_data_obj[key].pivot(index='Height',columns='DataStartTime',values='Value')
        fig, ax = plt.subplots()
        X, Y = np.meshgrid(np.arange(len(df.columns)), df.index)
        c = ax.contour(X, Y, df, levels=8, colors="black", linewidths=0.5)
        ax.clabel(c, inline=True, fontsize=10)
        c = ax.contourf(X, Y, df, levels=40, cmap='jet')
        plt.colorbar(c)
        ax.set_xlabel('时间')
        ax.set_ylabel('高度(km)')
        if key == "T":
            ax.set_title('温度时空分布(K)')
        elif key == "D":
            ax.set_title('大气密度时空分布(%)')
        elif key == "N":
            ax.set_title('钠密度时空分布(%)')
        elif key == "F":
            ax.set_title('铁密度时空分布(%)')
        elif key == "W1":
            ax.set_title('纬向风时空分布(m/s)')
        elif key == "W2":
            ax.set_title('经向风时空分布(m/s)')
        elif key == "W":
            ax.set_title('风速时空分布(m/s)')
        elif key == "X":
            ax.set_title('风向时空分布(°)')
        fig_path = os.path.join(figure_conf["figure_location"], key+'.png')
        fig.savefig(fig_path, dpi=300, bbox_inches="tight")
        img_dict[key] = fig_path
    return img_dict
